package com.xh.hongaicodegenerator.utils;

import cn.hutool.core.img.ImgUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import com.xh.hongaicodegenerator.exception.BusinessException;
import com.xh.hongaicodegenerator.exception.ErrorCode;
import io.github.bonigarcia.wdm.WebDriverManager;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.openqa.selenium.JavascriptExecutor;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeOptions;
import org.openqa.selenium.support.ui.WebDriverWait;

import java.io.File;
import java.time.Duration;
import java.util.UUID;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 截图工具类
 *
 */
@Slf4j
public class WebScreenShotUtils {
    private static final int MIN_DRIVER_COUNT = 2;   // 最小连接数
    private static final int MAX_DRIVER_COUNT = 4;  // 最大连接数
    // 在类中添加新的常量
    private static final int QUEUE_THRESHOLD = 8; // 排队任务阈值
    private static final long DRIVER_IDLE_TIMEOUT = 60; // 空闲 driver 超时销毁秒数
    // 添加新的成员变量来记录等待任务数
    private static final AtomicInteger waitingTaskCount = new AtomicInteger(0);


    // webDriver 池，并支持动态扩容，定时清理空闲 driver
    private static final BlockingQueue<WebDriver> driverPool = new LinkedBlockingQueue<>();
    private static final AtomicInteger currentDriverCount = new AtomicInteger(0);

    // 初始化最小连接数，最开始1个
    static {
        for (int i = 0; i < 1; i++) {
            driverPool.offer(initChromeDriver(1600, 900));
            currentDriverCount.incrementAndGet();
        }
        // 启动后台线程定时清理空闲 driver
        startDriverCleaner();
    }

    @PreDestroy
    public void destroy() {
        for (WebDriver driver : driverPool) {
            try {
                driver.quit();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        driverPool.clear();
    }


    /** 从池中取 driver（支持动态扩容） */
    private static WebDriver borrowDriver() throws InterruptedException {
        //线程安全，乐观锁
        WebDriver driver = driverPool.poll();
        if (driver != null) {
            return driver;
        }

        // 没有可用driver，增加等待任务计数
        waitingTaskCount.incrementAndGet();
        try {
            // 检查是否需要创建新的WebDriver
            if (currentDriverCount.get() < MAX_DRIVER_COUNT) {
                synchronized (WebScreenShotUtils.class) {
                    // 再次检查条件
                    if (currentDriverCount.get() < MAX_DRIVER_COUNT) {
                        // 如果当前连接数小于最小连接数，或者等待任务数超过阈值，则创建新的WebDriver
                        if (currentDriverCount.get() < MIN_DRIVER_COUNT ||
                                waitingTaskCount.get() > QUEUE_THRESHOLD) {
                            log.info("动态增加webDriver，当前连接数: {}, 等待任务数: {}",
                                    currentDriverCount.get(), waitingTaskCount.get());
                            driver = initChromeDriver(1600, 900);
                            currentDriverCount.incrementAndGet();
                            return driver;
                        }
                    }
                }
            }

            // 等待可用 driver
            return driverPool.take();
        } finally {
            // 减少等待任务计数
            waitingTaskCount.decrementAndGet();
        }
    }



    /** 用完放回池子 */
    private static void returnDriver(WebDriver driver) {
        driverPool.offer(driver);
    }

    /** 定时清理多余空闲 driver */
    /** 定时清理多余空闲 driver */
    private static void startDriverCleaner() {
        log.info("启动后台线程定时清理多余空闲 driver");
        Executors.newSingleThreadScheduledExecutor()
                .scheduleAtFixedRate(() -> {
                    // 只要总连接数大于最小连接数，且有空闲driver，就进行销毁
                    while (currentDriverCount.get() > MIN_DRIVER_COUNT && driverPool.size() > 0) {
                        log.info("目前webDriver池中有{}个webDriver，有{}个空闲webDriver，进行webDriver销毁",
                                currentDriverCount.get(), driverPool.size());
                        WebDriver driver = driverPool.poll();  // 获取空闲driver
                        if (driver != null) {
                            try {
                                driver.quit();  // 销毁driver
                            } catch (Exception e) {
                                e.printStackTrace();
                            }
                            currentDriverCount.decrementAndGet();  // 减少总连接数计数
                        }
                    }
                }, DRIVER_IDLE_TIMEOUT, DRIVER_IDLE_TIMEOUT, TimeUnit.SECONDS);
    }



    /**
     * 生成网页截图
     *
     * @param webUrl 要截图的网址
     * @return 压缩后的截图文件路径，失败返回 null
     */
    public static String saveWebPageScreenshot(String webUrl) {
        // 非空校验
        if (StrUtil.isBlank(webUrl)) {
            log.error("网页截图失败，url为空");
            return null;
        }
        WebDriver webDriver = null;
        // 创建临时目录
        try {
            String rootPath = System.getProperty("user.dir") + "/tmp/screenshots/" + UUID.randomUUID().toString().substring(0, 8);
            FileUtil.mkdir(rootPath);
            // 图片后缀
            final String IMAGE_SUFFIX = ".png";
            // 原始图片保存路径
            String imageSavePath = rootPath + File.separator + RandomUtil.randomNumbers(5) + IMAGE_SUFFIX;

            // 从池中获取 driver（如果没有可用的，会阻塞等待）
            webDriver = borrowDriver();;
            // 访问网页
            webDriver.get(webUrl);
            // 等待网页加载
            waitForPageLoad(webDriver);
            // 截图
            byte[] screenshotBytes = ((TakesScreenshot) webDriver).getScreenshotAs(OutputType.BYTES);
            // 保存原始图片
            saveImage(screenshotBytes, imageSavePath);
            log.info("原始截图保存成功：{}", imageSavePath);
            // 压缩图片
            final String COMPRESS_SUFFIX = "_compressed.jpg";
            String compressedImagePath = rootPath + File.separator + RandomUtil.randomNumbers(5) + COMPRESS_SUFFIX;
            compressImage(imageSavePath, compressedImagePath);
            log.info("压缩图片保存成功：{}", compressedImagePath);
            // 删除原始图片
            FileUtil.del(imageSavePath);
            return compressedImagePath;
        } catch (Exception e) {
            log.error("网页截图失败：{}", webUrl, e);
            return null;
        } finally {
            // 用完 driver 放回池子
            if (webDriver != null) {
                try {
                    returnDriver(webDriver);
                } catch (Exception e) {
                    log.warn("归还 WebDriver 池失败", e);
                }
            }
        }
    }

    /**
     * 初始化 Chrome 浏览器驱动
     */
    private static WebDriver initChromeDriver(int width, int height) {
        try {
            // 自动管理 ChromeDriver
            WebDriverManager.chromedriver().setup();
            // 配置 Chrome 选项
            ChromeOptions options = new ChromeOptions();
            // 无头模式
            options.addArguments("--headless=new");
            // 禁用GPU（在某些环境下避免问题）
            options.addArguments("--disable-gpu");
            // 禁用沙盒模式（Docker环境需要）
            options.addArguments("--no-sandbox");
            // 禁用开发者shm使用
            options.addArguments("--disable-dev-shm-usage");
            // 设置窗口大小
            options.addArguments(String.format("--window-size=%d,%d", width, height));
            // 禁用扩展
            options.addArguments("--disable-extensions");
            // 设置用户代理
            // 禁用图片、动画，降低内存消耗
            options.addArguments("--blink-settings=imagesEnabled=false");
            options.addArguments("--disable-animations");

            options.addArguments("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36");
            // 创建驱动
            WebDriver driver = new ChromeDriver(options);
            // 设置页面加载超时
            driver.manage().timeouts().pageLoadTimeout(Duration.ofSeconds(120));
            // 设置隐式等待
            driver.manage().timeouts().scriptTimeout(Duration.ofSeconds(60));    // JS 脚本超时
            driver.manage().timeouts().implicitlyWait(Duration.ofSeconds(10));
            return driver;
        } catch (Exception e) {
            log.error("初始化 Chrome 浏览器失败", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "初始化 Chrome 浏览器失败");
        }
    }

    /**
     * 保存图片到文件
     *
     * @param imageBytes
     * @param imagePath
     */
    private static void saveImage(byte[] imageBytes, String imagePath) {
        try {
            FileUtil.writeBytes(imageBytes, imagePath);
        } catch (Exception e) {
            log.error("保存图片失败：{}", imagePath, e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "保存图片失败");
        }
    }

    /**
     * 压缩图片
     *
     * @param originImagePath
     * @param compressedImagePath
     */
    private static void compressImage(String originImagePath, String compressedImagePath) {
        // 压缩图片质量（0.1 = 10% 质量）
        final float COMPRESSION_QUALITY = 0.5f;
        try {
            ImgUtil.compress(
                    FileUtil.file(originImagePath),
                    FileUtil.file(compressedImagePath),
                    COMPRESSION_QUALITY
            );
        } catch (Exception e) {
            log.error("压缩图片失败：{} -> {}", originImagePath, compressedImagePath, e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "压缩图片失败");
        }
    }

    /**
     * 等待页面加载完成
     *
     * @param webDriver
     */
    private static void waitForPageLoad(WebDriver webDriver) {
        try {
            // 创建等待页面加载对象
            WebDriverWait wait = new WebDriverWait(webDriver, Duration.ofSeconds(10));
            // 等待 document.readyState 为 complete
            wait.until(driver -> ((JavascriptExecutor) driver)
                    .executeScript("return document.readyState").
                    equals("complete")
            );
            // 额外等待一段时间，确保动态内容加载完成
            Thread.sleep(3000);
            log.info("页面加载完成");
        } catch (Exception e) {
            log.error("等待页面加载时出现异常，继续执行截图", e);
        }
    }
}

